import cv2
import sys
sys.path.append("../")
sys.path.append("./")

import torch.nn
import torch
import numpy as np
import DataLoader.SimuDataLoader as SimuDataLoader
import DataLoader.DataUtils as DataUtils


class EdgeColorLoss(torch.nn.Module):
    def __init__(self, colorWeight=0.3, edgeWeight=0.3, basicLossFunc=torch.nn.L1Loss()):
        super().__init__()
        self.colorWeight = colorWeight
        self.edgeWeight = edgeWeight
        self.basicLossFunc = basicLossFunc
    
    
    def edge(self, inputImg):
        # print(inputImg.shape)
        kernel = [[-1, -2, -1], [-2, 19, -2],[-1, -2, -1]] # LoG Kernel
        kernel = torch.FloatTensor(kernel).cuda().unsqueeze(0).unsqueeze(0) / 7
        weight = torch.nn.Parameter(data=kernel, requires_grad=False)
        
        r, g, b = inputImg[:,0:1,:], inputImg[:,1:2,:], inputImg[:,2:3,:]
        
        r_edge = torch.nn.functional.conv2d(r, weight, bias=None, stride=1, padding=1)
        g_edge = torch.nn.functional.conv2d(g, weight, bias=None, stride=1, padding=1)
        b_edge = torch.nn.functional.conv2d(b, weight, bias=None, stride=1, padding=1)
        
        return torch.cat([r_edge, g_edge, b_edge], dim=1)
    
    def edgeLoss(self, inputImg, target):
        inputEdge = self.edge(inputImg)
        targetEdge = self.edge(target) * 1.2
        
        return self.basicLossFunc(inputEdge, targetEdge), torch.abs(inputEdge-targetEdge)
    
    def colorLoss(self, inputImg, target):
        cosEbLoss = torch.nn.CosineEmbeddingLoss(reduction = "mean")
        return cosEbLoss(inputImg, target, torch.ones(1).to(inputImg.device))
    
    def forward(self, inputImg, target):
        basicLoss = self.basicLossFunc(inputImg, target)
        # punishment = self.calcPunishment(inputImg, target)
        colorLoss = self.colorLoss(inputImg, target)
        edgeLoss, edgeResidue = self.edgeLoss(inputImg, target)
        
        basicWeight = (1-self.colorWeight-self.edgeWeight)
        return (basicLoss * basicWeight + colorLoss * self.colorWeight + edgeLoss * self.edgeWeight, 
                basicLoss, colorLoss, edgeLoss, edgeResidue)
    
    

    
if __name__ == "__main__":
    pass